import matplotlib.pyplot as plt
import numpy as np
import torch
from tqdm import tqdm  # tqdm是显示循环进度条的库
import env


class QLearning:
    """ Q-learning算法 """

    def __init__(self, ncol, nrow, epsilon, alpha, gamma, n_action=21):
        self.Q_table = np.zeros([20, n_action])  # 初始化Q(s,a)表格
        self.n_action = n_action  # 动作个数
        self.alpha = alpha  # 学习率
        self.gamma = gamma  # 折扣因子
        self.epsilon = epsilon  # epsilon-贪婪策略中的参数

    def take_action(self, state):  # 选取下一步的操作
        # if np.random.random() < self.epsilon:
        action = np.random.randint(self.n_action)
        # else:
        #     action = np.argmax(self.Q_table[state])
        return action

    def best_action(self, state):  # 用于打印策略
        Q_max = np.max(self.Q_table[state])
        a = [0 for _ in range(self.n_action)]
        for i in range(self.n_action):
            if self.Q_table[state, i] == Q_max:
                a[i] = 1
        return a

    def update(self, s0, a0, r, s1):
        """
        核心策略
        :param s0:
        :param a0:
        :param r:
        :param s1:
        :return:
        """
        td_error = r + self.gamma * self.Q_table[s1].max(
        ) - self.Q_table[s0, a0]
        self.Q_table[s0, a0] += self.alpha * td_error


# 需要修改
ncol = 12
nrow = 4
# 需要修改
env = env.Env()
np.random.seed(0)
epsilon = 0.1
alpha = 0.1
gamma = 0.9
agent = QLearning(ncol, nrow, epsilon, alpha, gamma)
num_episodes = 500  # 智能体在环境中运行的序列的数量

return_list = []  # 记录每一条序列的回报
for i in range(10):  # 显示10个进度条
    # tqdm的进度条功能
    with tqdm(total=int(num_episodes / 10), desc='Iteration %d' % i) as pbar:
        for i_episode in range(int(num_episodes / 10)):  # 每个进度条的序列数
            episode_return = 0
            cur_cap, mask = env.reset()
            selected_vehicle = torch.tensor([[2], [0], [2], [0], [1]])
            current_node = torch.tensor([[15], [8], [16], [2], [13]])
            # 根据当前状态选择action
            action = agent.take_action(mask)
            # 通过action得到下一步state,reward,done
            mask, next_state, reward_all = env.step(selected_vehicle, current_node)
            episode_return += reward_all  # 这里回报的计算不进行折扣因子衰减
            agent.update(state, action, mask, next_state)
            state = next_state
            return_list.append(episode_return)
            if (i_episode + 1) % 10 == 0:  # 每10条序列打印一下这10条序列的平均回报
                pbar.set_postfix({
                    'episode':
                        '%d' % (num_episodes / 10 * i + i_episode + 1),
                    'return':
                        '%.3f' % np.mean(return_list[-10:])
                })
            pbar.update(1)

episodes_list = list(range(len(return_list)))
plt.plot(episodes_list, return_list)
plt.xlabel('Episodes')
plt.ylabel('Returns')
plt.title('Q-learning on {}'.format('Cliff Walking'))
plt.show()
